import os
import shutil
from . import hogfeature
from . import knnclassify, svmclassify, rfclassify
from sklearn.svm import SVC
import numpy as np
import glob
import time
import joblib
from . import variable as va  

'''
=================
    支持方法
=================

'''
# 加载数据集特征文件
def load_feats(filePath):
    features = []
    labels = []
    # 加载所有的特征文件
    for feat_path in glob.glob(os.path.join(filePath, '*.feat')):
        data = joblib.load(feat_path)
        # 每个特征文件中都有一个样本，最后一个数据是label，其余的是特征
        features.append(data[:-1])
        labels.append(data[-1])
    # 返回特征与标签的列表
    return features, labels 


# 把模型测试时的结果写入到文件中
def write_to_txt(list, name):
    f = open(va.result_path + name + '_result.txt', 'w')
    f.writelines(list)
    f.close()


# 保存模型的设置与性能
def save_model_cfg(clf_type, args, result):
    f = open(va.model_path + clf_type + '_cfg.txt', 'w')
    for key in args.keys():
        f.write(key + '\t'+ str(args[key]) + '\n')
    f.write('训练图片总数：\t' + str(result['train_num']) + '\n')
    f.write('训练时间：\t' + str(result['train_time']) + '\n')
    f.write('测试时正确分类的数量：\t' + str(result['correct']) + '\n')
    f.write('测试图片总数：\t' + str(result['test_num']) + '\n')
    f.write('测试时的正确率：\t' + ('%.2f' % result['rate']) + '\n')
    f.write('测试时间：\t' + str(result['test_time']) + '\n')
    f.close()
    

# 检查用于存放特征的文件夹是否存在，如果有就先删掉--重建
def remkdir():
    # shutil.rmtree()函数能递归删除文件夹下的所有子文件夹和子文件。
    if os.path.exists(va.train_feat_path):
        shutil.rmtree(va.train_feat_path)
    if os.path.exists(va.test_feat_path):
        shutil.rmtree(va.test_feat_path)
    # 创建文件夹
    os.mkdir(va.train_feat_path)
    os.mkdir(va.test_feat_path)


# 创建分类器
def buildClassify(clf_type, args):
    clf = None
    if clf_type == 'KNN':
        clf = knnclassify.buildClassify(args)
    elif clf_type == 'SVM':
        clf = svmclassify.buildClassify(args)
    elif clf_type == 'RF':
        clf = rfclassify.buildClassify(args)
    return clf


# 训练
def train(clf, name):
    # 加载训练集的特征与标签
    features, labels = load_feats(va.train_feat_path)
    t0 = time.time()
    clf.fit(features, labels)
    
    # 如果保存模型的文件夹不存在，创建这个文件夹
    if not os.path.exists(va.model_path):
        os.makedirs(va.model_path)
    # 保存模型
    joblib.dump(clf, va.model_path + name + '_model')
    t1 = time.time()
    train_time = t1 - t0
    return clf, train_time, len(labels)

# 测试--返回分类正确的个数、测试总数与测试时间
def test(clf, name):
    t0 = time.time()
    correct_number = 0 # 分类正确的次数
    total = 0          # 分类总次数
    result_list = []   # 分类结果列表
    for feat_path in glob.glob(os.path.join(va.test_feat_path, '*.feat')):
        total += 1
        symbol = '\\'
        # 获取测试集中图像的特征
        image_name = feat_path.split(symbol)[1].split('.feat')[0]
        data_test = joblib.load(feat_path)

        # 规范化label
        data_test_feat = data_test[:-1].reshape((1, -1)).astype(np.float64)
        # 预测结果 -- 会得到一个浮点数
        result = clf.predict(data_test_feat)
        # 记录该图片的预测结果
        result_list.append(image_name + ' --> ' + va.labels[int(result[0])] + '\n')
        if int(result[0]) == data_test[-1]:
           correct_number += 1
        else:
            result_list[-1] = '【Error】 ' + result_list[-1]
    write_to_txt(result_list, name)
    rate = float(correct_number) / total
    t1 = time.time()
    test_time = t1 - t0
    result = {}
    result['correct'] = correct_number
    result['test_num'] = total
    result['rate'] = rate
    result['test_time'] = test_time 
    return result

    
'''
===============
    API函数
===============
'''
# 训练模（指定模型类型 + 模型参数）
def train_model(clf_type, args):
    # 创建分类器
    clf = buildClassify(clf_type, args)
    # 重新创建用于存放特征的文件夹
    remkdir()
    # 获取数据集的特征
    hogfeature.get_feats()
    # 训练模型
    model, train_time, train_num = train(clf, clf_type)
    # 进行测试
    result = test(model, clf_type)
    # 组装结果
    result['train_time'] = train_time
    result['train_num'] = train_num
    # 保存模型的配置与性能
    save_model_cfg(clf_type, args, result)
    return result
